package com.jwiki.app.jkeeper.common.jdbc;

import java.io.StringReader;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.Collections;
import java.util.Properties;

import org.apache.commons.lang3.ArrayUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.ibatis.executor.parameter.ParameterHandler;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Plugin;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.reflection.SystemMetaObject;
import org.apache.ibatis.scripting.defaults.DefaultParameterHandler;
import org.apache.ibatis.session.RowBounds;
import org.springframework.stereotype.Component;

import net.sf.jsqlparser.JSQLParserException;
import net.sf.jsqlparser.parser.CCJSqlParserManager;
import net.sf.jsqlparser.statement.select.PlainSelect;
import net.sf.jsqlparser.statement.select.Select;
import net.sf.jsqlparser.statement.select.SelectBody;

@Component
@SuppressWarnings({"rawtypes", "unchecked"})
@Intercepts({@Signature(type=StatementHandler.class, method="prepare", args={Connection.class, Integer.class})})
public class JdbcUtils implements Interceptor {

    private static final String TARGET_GETTER_KEY = "target";

    private static CCJSqlParserManager sqlParserManager = new CCJSqlParserManager();

    private static final ThreadLocal<Parameter> PARAMETER = new ThreadLocal<>();

    public static class Builder {

        private Parameter parameter;

        public Builder(Parameter parameter) {
            this.parameter = parameter;
        }

        public Builder groupBy(String... columns) {
            parameter.setGroupBy(columns);
            return this;
        }

        public Builder orderBy(String... columns) {
            parameter.setOrderBy(columns);
            return this;
        }

        public Builder limit(int limit) {
            return this.limit(0, limit);
        }

        public Builder limit(int offset, int limit) {
            Page page = new Page(offset, limit);
            parameter.setPage(page);
            return this;
        }

        public <E> Page<E> getPage() {
            return parameter.getPage();
        }

    }

    public static Builder build() {
        Parameter parameter = new Parameter();
        PARAMETER.set(parameter);
        return new Builder(parameter);
    }

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        try {
            Parameter param = PARAMETER.get();
            if (param != null) {
                Object target = invocation.getTarget();
                if (target instanceof StatementHandler) {
                    StatementHandler statement = (StatementHandler) target;
                    BoundSql boundSql = getSystemMetaObject(statement, "delegate.boundSql");
                    StringBuilder sql = new StringBuilder(boundSql.getSql());

                    if (ArrayUtils.isNotEmpty(param.getGroupBy())) {
                        sql.append(" GROUP BY ");
                        sql.append(StringUtils.join(param.getGroupBy(), ", "));
                    }
                    if (ArrayUtils.isNotEmpty(param.getOrderBy())) {
                        sql.append(" ORDER BY ");
                        sql.append(StringUtils.join(param.getOrderBy(), ", "));
                    }

                    Page page = param.getPage();
                    if (page != null) {
                        MappedStatement mappedStatement = getSystemMetaObject(statement, "delegate.mappedStatement");
                        Connection connection = (Connection) invocation.getArgs()[0];
                        // select count sql
                        String countSql = getSimpleCountSql(sql.toString());

                        try (PreparedStatement countStatement = connection.prepareStatement(countSql)) {
                            BoundSql countBoundSql = new BoundSql(mappedStatement.getConfiguration(), countSql,
                                    boundSql.getParameterMappings(), boundSql.getParameterObject());
                            ParameterHandler parameterHandler = new DefaultParameterHandler(mappedStatement,
                                    boundSql.getParameterObject(), countBoundSql);
                            parameterHandler.setParameters(countStatement);

                            ResultSet rsCount = countStatement.executeQuery();
                            long total = 0;
                            if (rsCount.next()) {
                                total = rsCount.getLong(1);
                            }
                            page.setTotal(total);
                        } catch (SQLException e) {
                            throw e;
                        }

                        setSystemMetaObject(statement, "delegate.rowBounds.offset", RowBounds.NO_ROW_OFFSET);
                        setSystemMetaObject(statement, "delegate.rowBounds.limit", RowBounds.NO_ROW_LIMIT);

                        sql.append(" LIMIT ");
                        sql.append(page.getOffset());
                        sql.append(", ");
                        sql.append(page.getLimit());
                    }

                    setSystemMetaObject(statement, "delegate.boundSql.sql", sql.toString());
                }
            }
            return invocation.proceed();
        } catch (Exception e) {
            throw e;
        } finally {
            PARAMETER.remove();
        }
    }

    protected String getSimpleCountSql(String sql) throws JSQLParserException {
        SelectBody selectBody = ((Select) sqlParserManager.parse(new StringReader(sql))).getSelectBody();
        if (selectBody instanceof PlainSelect) {
            PlainSelect select = (PlainSelect) selectBody;
            select.setOrderByElements(null);
            if (select.getGroupByColumnReferences() == null) {
                select.setSelectItems(Collections.singletonList("COUNT(1)"));
                return select.toString();
            }
        }
        return StringUtils.join("SELECT COUNT(1) FROM (", sql, ") _TEMP_");
    }

    private <T> T getSystemMetaObject(Object mataObject, String metaKey) {
        MetaObject metaStatement = SystemMetaObject.forObject(mataObject);
        while (metaStatement.hasGetter(TARGET_GETTER_KEY)) {
            Object object = metaStatement.getValue(TARGET_GETTER_KEY);
            metaStatement = SystemMetaObject.forObject(object);
        }
        return (T) metaStatement.getValue(metaKey);
    }

    private void setSystemMetaObject(Object mataObject, String metaKey, Object metaValue) {
        MetaObject metaStatement = SystemMetaObject.forObject(mataObject);
        while (metaStatement.hasGetter(TARGET_GETTER_KEY)) {
            Object object = metaStatement.getValue(TARGET_GETTER_KEY);
            metaStatement = SystemMetaObject.forObject(object);
        }
        metaStatement.setValue(metaKey, metaValue);
    }

    @Override
    public Object plugin(Object target) {
        if (target instanceof StatementHandler) {
            return Plugin.wrap(target, this);
        } else {
            return target;
        }
    }

    @Override
    public void setProperties(Properties properties) {
        // nothing
    }

}
